﻿#include <iomanip>
#include <sstream>

#include "coarsening.hpp"
#include "solverFactory.hpp"
#include "unapMultigrid.hpp"

#ifdef SW_SLAVE
#include "vectorOps.h"
#endif

#ifdef TIMERS
#include "utilityTimers.hpp"
#endif

template <typename Cmpt>
UNAP::MGSolver<Cmpt>::MGSolver(const Matrix<Cmpt> &A,
                               const Coarsening<Cmpt> &MGAgglomeration)
    : Solver<Cmpt>(A.getCommunicator()),
      finestMatrix_(A),
      nPreSweeps_(2),
      nPostSweeps_(2),
      maxPreSweeps_(2),
      maxPostSweeps_(2),
      nFinestSweeps_(2),
      preSweepsLevelMultiplier_(1),
      postSweepsLevelMultiplier_(1),
      scaleCorrection_(A.symm()),
      coarsening_(&MGAgglomeration),
      smootherType_("Jacobi"),
      Smoothers_(new PtrList<Smoother<Cmpt> >),
      hasSmoothers_(true),
      coarsestPrecond_(A.symm() ? "DIC" : "DILU"),
      coarsestSolver_(A.symm() ? "CG" : "BiCGStab"),
      levels_(MGAgglomeration.size() + 1),
      noCoarse_(MGAgglomeration.ifNoCoarse()),
      levelKcycle_(2) {}

template <typename Cmpt>
UNAP::MGSolver<Cmpt>::MGSolver(const Matrix<Cmpt> &A,
                               const Coarsening<Cmpt> &MGAgglomeration,
                               PtrList<Smoother<Cmpt> > &Smoothers)
    : Solver<Cmpt>(A.getCommunicator()),
      finestMatrix_(A),
      nPreSweeps_(2),
      nPostSweeps_(2),
      maxPreSweeps_(2),
      maxPostSweeps_(2),
      nFinestSweeps_(2),
      preSweepsLevelMultiplier_(1),
      postSweepsLevelMultiplier_(1),
      scaleCorrection_(A.symm()),

      coarsening_(&MGAgglomeration),
      smootherType_("Jacobi"),
      Smoothers_(&Smoothers),
      hasSmoothers_(false),
      coarsestPrecond_(A.symm() ? "DIC" : "DILU"),
      coarsestSolver_(A.symm() ? "CG" : "BiCGStab"),
      levels_(MGAgglomeration.size() + 1),
      noCoarse_(MGAgglomeration.ifNoCoarse()),
      levelKcycle_(2) {}

template <typename Cmpt>
UNAP::MGSolver<Cmpt>::~MGSolver() {
  if (hasSmoothers_) {
    DELETE_POINTER(Smoothers_)
  }
}

template <typename Cmpt>
Cmpt UNAP::MGSolver<Cmpt>::scalingFactor(Vector<Cmpt> &field,
                                         const Vector<Cmpt> &source,
                                         const Matrix<Cmpt> &A) const {
  Vector<Cmpt> Acf(field);
  A.spMV(Acf, field);

  Cmpt sfVal[2] = {0.0, 0.0};
  Vector<Cmpt> &D = A.diag();
  Cmpt *DPtr = D.data();
  Cmpt *fieldPtr = field.data();
  Cmpt *AcfPtr = Acf.data();
  const Cmpt *sourcePtr = source.data();
  label rowSize = field.size();

  IFNOT_SWACC_SMALL(rowSize) {
    forAll(i, rowSize) {
      sfVal[0] += sourcePtr[i] * fieldPtr[i];
      sfVal[1] += AcfPtr[i] * fieldPtr[i];
      fieldPtr[i] += (sourcePtr[i] - AcfPtr[i]) / DPtr[i];
    }
  }
#ifdef SW_SLAVE
  else {
    MVM_Arrays arrays1;
    init_MVM_Arrays(&arrays1, rowSize);
    arrays1.A1Ptr = fieldPtr;
    arrays1.A2Ptr = (scalar *)sourcePtr;
    arrays1.A3Ptr = AcfPtr;
    arrays1.A4Ptr = DPtr;
    arrays1.k1Ptr = &sfVal[0];
    arrays1.k2Ptr = &sfVal[1];
    scalingFactor_host(&arrays1);
  }
#endif

  reduceSum(&sfVal[0], this->commcator_, 2);
  return sfVal[0] / (sfVal[1] + SMALL);
}

template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::solveCoarsestLevel(
    Vector<Cmpt> &coarsestCorrField, const Vector<Cmpt> &coarsestSource) const {
  scalar relTol = 1.0e-8;
  scalar tol = 0.0;
  label maxIt = 100;
  bool print = false;
  SolverPerformance coarseSolverPerf;

  const Matrix<Cmpt> &mat = matrix(levels_ - 1);
  Preconditioner<Cmpt> *precond =
      PrecondFactory::buildPrecond((Matrix<Cmpt> &)mat, coarsestPrecond_);
  Solver<Cmpt> *solver = SolverFactory::buildSolver(*precond, coarsestSolver_);

  solver->SET_tolerance(tol);
  solver->SET_relTol(relTol);
  solver->SET_maxIter(maxIt);
  solver->SET_ifPrint(print);
  coarseSolverPerf = solver->solve(coarsestCorrField, mat, coarsestSource);
  delete precond;
  delete solver;
}

template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::SET_nPreSweeps(const label nPreSweeps) {
  nPreSweeps_ = nPreSweeps;
}
template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::SET_nPostSweeps(const label nPostSweeps) {
  nPostSweeps_ = nPostSweeps;
}
template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::SET_nFinestSweeps(const label nFinestSweeps) {
  nFinestSweeps_ = nFinestSweeps;
}
template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::SET_maxPreSweeps(const label maxPreSweeps) {
  maxPreSweeps_ = maxPreSweeps;
}
template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::SET_maxPostSweeps(const label maxPostSweeps) {
  maxPostSweeps_ = maxPostSweeps;
}
template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::SET_preSweepsLevelMultiplier(
    const label multiplier) {
  preSweepsLevelMultiplier_ = multiplier;
}
template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::SET_postSweepsLevelMultiplier(
    const label multiplier) {
  postSweepsLevelMultiplier_ = multiplier;
}
template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::SET_scaleCorrection(bool scaleCorrection) {
  scaleCorrection_ = scaleCorrection;
}

template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::SET_smootherType(string type) {
  SET_smootherType(type.c_str());
}
template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::SET_smootherType(const char *type) {
  smootherType_ = type;
}

template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::initSmoothers(PtrList<Smoother<Cmpt> > *smPtr) {
  if (smPtr != NULL) {
    if (hasSmoothers_) {
      DELETE_POINTER(Smoothers_)
    }
    Smoothers_ = smPtr;
    return;
  }
  Smoothers_->SET_size(levels_);
  const string type(smootherType_);
  forAll(i, levels_) {
    Smoother<Cmpt> *smoothi =
        SmootherFactory::buildSmoother(this->commcator_, type);
    Smoothers_->setLevel(i, *smoothi);
    (*Smoothers_)[i].init();
  }
}

template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::setUp() {}

template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::SET_coarsestPrecond(const string &precondName) {
  SET_coarsestPrecond(precondName.c_str());
}
template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::SET_coarsestPrecond(const char *precondName) {
  coarsestPrecond_ = precondName;
}
template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::SET_coarsestSolver(const string &solverName) {
  SET_coarsestSolver(solverName.c_str());
}
template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::SET_coarsestSolver(const char *solverName) {
  coarsestSolver_ = solverName;
  if (std::strcmp(solverName, "Direct") == 0) SET_coarsestPrecond("LU");
}

template <typename Cmpt>
template <typename T>
void UNAP::MGSolver<Cmpt>::restrictField(Vector<T> &cf, const Vector<T> &ff,
                                         const label fineLevelIndex) const {
  coarsening_->restrictField(cf, ff, fineLevelIndex);
}

template <typename Cmpt>
template <typename T>
void UNAP::MGSolver<Cmpt>::prolongField(Vector<T> &ff, const Vector<T> &cf,
                                        const label coarseLevelIndex) const {
  coarsening_->prolongField(ff, cf, coarseLevelIndex);
}

template <typename Cmpt>
const UNAP::Matrix<Cmpt> &UNAP::MGSolver<Cmpt>::matrix(
    const label leveli) const {
  return coarsening_->matrixLevel(leveli);
}

template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::printStat(bool ifPrintProc) {
  std::cout.flags(std::ios::right);
  std::ostream *file = NULL;

  forAll(i, 38) COUT << "=";
  COUT << ENDL;
  std::cout.width(3);
  COUT << "lev";
  std::cout.width(13);
  COUT << "rows";
  std::cout.width(13);
  COUT << "nnz";
  std::cout.width(9);
  COUT << "sparse" << ENDL;
  forAll(i, 38) COUT << "=";
  COUT << ENDL;

  if (ifPrintProc) {
    file = POUT.stream();

    forAll(i, 76) POUT << "=";
    POUT << ENDL;
    file->flags(std::ios::right);
    file->width(3);
    *file << "lev";
    file->width(13);
    *file << "lrows"
          << "/";
    file->flags(std::ios::left);
    file->width(13);
    *file << "rows";
    file->flags(std::ios::right);
    file->width(13);
    *file << "lnnz"
          << "/";
    file->flags(std::ios::left);
    file->width(13);
    *file << "nnz";
    file->flags(std::ios::right);
    file->width(12);
    *file << "lsparse"
          << "/";
    file->flags(std::ios::left);
    file->width(6);
    *file << "sparse" << ENDL;
    forAll(i, 76) POUT << "=";
    POUT << ENDL;
  }

  scalar opComplexity = 0.;
  scalar gridComplexity = 0.;
  scalar fineGridSize, fineNnz;

  forAll(li, levels_) {
    const Matrix<Cmpt> &Ai = matrix(li);

    const label rowSize = Ai.size();
    label totalSize = rowSize;
    const label nnz = Ai.lnnz();
    label64 totalnnz = nnz;
    reduceSum(&totalSize, Ai.getCommunicator());
    reduceSum(&totalnnz, Ai.getCommunicator());
    if (li == 0) {
      fineGridSize = totalSize;
      fineNnz = totalnnz;
    }
    opComplexity += totalnnz / fineNnz;
    gridComplexity += totalSize / fineGridSize;
    const scalar lsparse = 1. * nnz / rowSize / totalSize;
    const scalar sparse = 1. * totalnnz / totalSize / totalSize;

    if (ifPrintProc) {
      file->flags(std::ios::right);
      file->flags(std::ios::fixed);
      file->precision(4);
      file->width(3);
      *file << li;
      file->width(13);
      *file << rowSize << "/";
      file->flags(std::ios::left);
      file->width(13);
      *file << totalSize;
      file->flags(std::ios::right);
      file->width(13);
      *file << nnz << "/";
      file->flags(std::ios::left);
      file->width(13);
      *file << totalnnz;
      file->flags(std::ios::right);
      file->flags(std::ios::fixed);
      file->precision(4);
      file->width(12);
      *file << lsparse << "/" << sparse << ENDL;
    }
    std::cout.flags(std::ios::right);
    std::cout.flags(std::ios::fixed);
    std::cout.precision(4);
    std::cout.width(3);
    COUT << li;
    std::cout.width(13);
    COUT << totalSize;
    std::cout.width(13);
    COUT << totalnnz;
    std::cout.width(9);
    COUT << sparse << ENDL;
  }
  forAll(i, 38) COUT << "=";
  COUT << ENDL;
  if (ifPrintProc) {
    forAll(i, 76) POUT << "=";
    POUT << ENDL;
  }

  std::cout.flags(std::ios::left);
  std::cout.flags(std::ios::fixed);
  std::cout.precision(8);
  std::cout.width(21);
  COUT << "Grid complexity: ";
  std::cout.flags(std::ios::right);
  std::cout.width(15);
  COUT << gridComplexity << ENDL;
  std::cout.flags(std::ios::left);
  std::cout.width(21);
  COUT << "Operator complexity: ";
  std::cout.flags(std::ios::right);
  std::cout.width(15);
  COUT << opComplexity << ENDL;
  forAll(i, 38) COUT << "=";
  COUT << ENDL;
}

template <typename Cmpt>
UNAP::SolverPerformance UNAP::MGSolver<Cmpt>::solve(Vector<Cmpt> &x,
                                                    const Matrix<Cmpt> &A,
                                                    const Vector<Cmpt> &b) {
  SolverPerformance solverPerf;
  const label rowSize = x.size();

  Vector<Cmpt> Apsi(rowSize, this->commcator_);

  Vector<Cmpt> finestCorrection(rowSize, this->commcator_);
  Vector<Cmpt> finestResidual(rowSize, this->commcator_);

  PtrList<Vector<Cmpt> > coarseCorrFields(levels_);

  PtrList<Vector<Cmpt> > coarseSources(levels_);

  forAll(leveli, levels_) {
    const label rowSize = matrix(leveli).size();
    Vector<Cmpt> *coarseCorrFieldsLeveliPtr =
        new Vector<Cmpt>(rowSize, this->commcator_);
    Vector<Cmpt> *coarseSourcesLeveliPtr =
        new Vector<Cmpt>(rowSize, this->commcator_);
    coarseCorrFields.setLevel(leveli, *coarseCorrFieldsLeveliPtr);
    coarseSources.setLevel(leveli, *coarseSourcesLeveliPtr);
  }

  A.spMV(Apsi, x);

  waxpy(finestResidual, -1., Apsi, b);

  solverPerf.initialResidual(this->normFactor(b));
#ifdef VERBOSE
  if (this->ifPrint_) solverPerf.printResidual();
#endif

  if (solverPerf.checkConvergence(this->tolerance_, this->relTol_,
                                  this->minIter_))
    return solverPerf;

#ifdef BEST_RESULT
  scalar bestResidual = solverPerf.initialResidual();
  Vector<Cmpt> bestSolve(x);
#endif
  do {
    solverPerf.nIterations()++;

    if (levelKcycle_ < 0) {
      Vcycle(0, x, b, coarseCorrFields, coarseSources, 1);
    } else {
      Kcycle(0, x, b, coarseCorrFields, coarseSources, 1);
    }

    A.spMV(Apsi, x);

    waxpy(finestResidual, -1., Apsi, b);
    solverPerf.finalResidual(this->normFactor(finestResidual));

#ifdef BEST_RESULT
    if (bestResidual > solverPerf.finalResidual()) {
      bestSolve = x;
      bestResidual = solverPerf.finalResidual();
    }
#endif

#ifdef VERBOSE
    if (this->ifPrint_) solverPerf.printResidual();
#endif
  } while ((solverPerf.nIterations() < this->maxIter_ &&
            !(solverPerf.checkConvergence(this->tolerance_, this->relTol_,
                                          this->minIter_))));

#ifdef BEST_RESULT
  x = bestSolve;
  if (ifPrint_) {
    std::cout.flags(std::ios::left);
    std::cout.flags(std::ios::showpoint);
    COUT << "At last";
    std::cout.width(5);
    COUT << this->nIterations() << " , res = ";
    std::cout.width(precision + 2);
    std::cout.precision(precision);
    COUT << bestResidual
         << " , rel res = " << bestResidual / this->initialResidual() << ENDL;
  }
#endif

  return solverPerf;
}

template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::precond(Vector<Cmpt> &w, const Vector<Cmpt> &r,
                                   label nCycle) {
  const label rowSize = r.size();
  Vector<Cmpt> Apsi(rowSize, this->commcator_);
  w.SET_zero();

  PtrList<Vector<Cmpt> > coarseCorrFields(levels_);

  PtrList<Vector<Cmpt> > coarseSources(levels_);
  forAll(leveli, levels_) {
    const label rowSize = matrix(leveli).size();
    Vector<Cmpt> *coarseCorrFieldsLeveliPtr =
        new Vector<Cmpt>(rowSize, this->commcator_);
    Vector<Cmpt> *coarseSourcesLeveliPtr =
        new Vector<Cmpt>(rowSize, this->commcator_);
    coarseCorrFields.setLevel(leveli, *coarseCorrFieldsLeveliPtr);
    coarseSources.setLevel(leveli, *coarseSourcesLeveliPtr);
  }

  for (label cycle = 0; cycle < nCycle; cycle++) {
    if (levelKcycle_ < 0) {
      Vcycle(0, w, r, coarseCorrFields, coarseSources, 1);
    } else {
      Kcycle(0, w, r, coarseCorrFields, coarseSources, 1);
    }
  }
}

template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::Vcycle(label level, Vector<Cmpt> &psi,
                                  const Vector<Cmpt> &source,
                                  PtrList<Vector<Cmpt> > &coarseCorrFields,
                                  PtrList<Vector<Cmpt> > &coarseSources,
                                  label k) {
  const label coarsestLevel = levels_ - 1;
  if (level == coarsestLevel) {
    solveCoarsestLevel(psi, source);

  } else {
    const label rowSize = coarseCorrFields[level].size();
    Vector<Cmpt> rA(rowSize, this->commcator_);
    Vector<Cmpt> vA(rowSize, this->commcator_);
    Vector<Cmpt> &coarseCorrFieldii = coarseCorrFields[level + 1];
    Vector<Cmpt> &coarseSourceii = coarseSources[level + 1];
    const Matrix<Cmpt> &mat = matrix(level);

    (*Smoothers_)[level].smooth(
        psi, mat, source,
        MIN(nPreSweeps_ + preSweepsLevelMultiplier_ * level, maxPreSweeps_));

    mat.spMV(rA, psi);
    aypx(rA, -1., source);

    restrictField(coarseSourceii, rA, level);

    coarseCorrFieldii = 0;

    Vcycle(level + 1, coarseCorrFieldii, coarseSourceii, coarseCorrFields,
           coarseSources, k > 1 ? (k - 1) : k);

    prolongField(vA, coarseCorrFields[level + 1], level);

    if (scaleCorrection_) vA.scale(scalingFactor(vA, rA, mat));

    axpy(psi, 1., vA);

    (*Smoothers_)[level].smooth(
        psi, mat, source,
        MIN(nPostSweeps_ + postSweepsLevelMultiplier_ * level, maxPostSweeps_));
  }
}

template <typename Cmpt>
void UNAP::MGSolver<Cmpt>::Kcycle(label level, Vector<Cmpt> &psi,
                                  const Vector<Cmpt> &source,
                                  PtrList<Vector<Cmpt> > &coarseCorrFields,
                                  PtrList<Vector<Cmpt> > &coarseSources,
                                  label k) {
  const label coarsestLevel = levels_ - 1;
  if (level == coarsestLevel)

    solveCoarsestLevel(psi, source);
  else {
    const label rowSize = coarseCorrFields[level].size();
    Vector<Cmpt> rA(rowSize, this->commcator_);
    Vector<Cmpt> vA(rowSize, this->commcator_);
    Vector<Cmpt> &coarseCorrFieldii = coarseCorrFields[level + 1];
    Vector<Cmpt> &coarseSourceii = coarseSources[level + 1];
    const Matrix<Cmpt> &mat = matrix(level);

    (*Smoothers_)[level].smooth(
        psi, mat, source,
        MIN(nPreSweeps_ + preSweepsLevelMultiplier_ * level, maxPreSweeps_));

    mat.spMV(rA, psi);
    aypx(rA, -1., source);

    restrictField(coarseSourceii, rA, level);

    const Matrix<Cmpt> &A = matrix(level + 1);
    Vector<Cmpt> ck(coarseCorrFieldii), rk(coarseSourceii), vk(ck), wk(ck),
        dk(ck);
    Cmpt rho1, alpha1, gamma, beta, rho2, alpha2;
    const scalar tFactor = 0.25;

    if (level > levelKcycle_) {
      Vcycle(level + 1, coarseCorrFieldii, coarseSourceii, coarseCorrFields,
             coarseSources, k);
    } else {
      Kcycle(level + 1, coarseCorrFieldii, coarseSourceii, coarseCorrFields,
             coarseSources, k);
    }
    ck = coarseCorrFieldii;
    A.spMV(vk, ck);
    if (A.symm()) {
      rho1 = dot(ck, vk);
      alpha1 = dot(ck, rk);
    } else {
      rho1 = vk.SumSqr();
      alpha1 = dot(vk, rk);
    }
    waxpy(coarseSourceii, -alpha1 / rho1, vk, rk);
    if (coarseSourceii.SumSqrt() / rk.SumSqrt() <= tFactor) {
      ck.scale(alpha1 / rho1);
      coarseCorrFieldii = ck;
    } else {
      rk = coarseSourceii;
      if (level > levelKcycle_) {
        Vcycle(level + 1, coarseCorrFieldii, coarseSourceii, coarseCorrFields,
               coarseSources, k);
      } else {
        Kcycle(level + 1, coarseCorrFieldii, coarseSourceii, coarseCorrFields,
               coarseSources, k);
      }
      dk = coarseCorrFieldii;
      A.spMV(wk, dk);
      if (A.symm()) {
        gamma = dot(dk, vk);
        beta = dot(dk, wk);
        alpha2 = dot(dk, rk);
      } else {
        gamma = dot(wk, vk);
        beta = wk.SumSqr();
        alpha2 = dot(wk, rk);
      }
      rho2 = beta - gamma * gamma / rho1;
      ck.scale(alpha1 / rho1 - gamma * alpha2 / rho1 / rho2);
      waxpy(coarseCorrFieldii, alpha2 / rho2, dk, ck);
    }

    prolongField(vA, coarseCorrFieldii, level);

    if (scaleCorrection_) vA.scale(scalingFactor(vA, rA, mat));

    axpy(psi, 1., vA);

    (*Smoothers_)[level].smooth(
        psi, mat, source,
        MIN(nPostSweeps_ + postSweepsLevelMultiplier_ * level, maxPostSweeps_));
  }
}

template class UNAP::MGSolver<scalar>;
